import os

import torch
from torch import optim
from torch.nn.utils import prune
from models.lwCETModel import lwCET
from myutils.modelUtils import get_configs

# 迭代式剪枝：不要一次性剪枝太多。更稳健的方法是采用迭代剪枝，即每次剪掉一小部分（如10%），然后微调，再剪下一部分，如此循环，直到达到目标稀疏度，这样能更好地保持模型性能

# 1. 加载模型
exp_log_dir = ''
test_model_name = os.path.join(exp_log_dir, "checkpoint_best.pt")
dataset_configs, hparams_class1 = get_configs()
hparams = hparams_class1.train_params
# 创建模型实例
model = lwCET(dataset_configs, hparams)
# model = ecgTransForm(configs=self.dataset_configs, hparams=self.hparams)
device = map_location = 'cuda:0' if torch.cuda.is_available() else 'cpu'
chkpoint = torch.load(test_model_name, device)
model.load_state_dict(chkpoint['model'])
model = model.to(device)

# 2 执行剪枝操作
parameters_to_prune = (
    (model.conv1, 'weight'),
    (model.conv2, 'weight'),
    (model.conv3, 'weight'),
)

# 3 应用全局非结构化剪枝
# 使用 l1_unstructured 方法，将 30% 的权重置零
# 结构化剪枝 prune.ln_structured（L2范数结构化剪枝），它直接移除整个神经元或通道，对硬件更友好
prune.global_unstructured(
    parameters_to_prune,
    pruning_method=prune.ln_structured,
    amount=0.3,  # 剪枝比例 30%
)

print("剪枝操作完成。")
# 查看剪枝效果：例如，第一个全连接层有多少权重变成了0（即被剪枝）
print(f"FC1层权重稀疏度: {(model.fc1.weight == 0).sum().item() / model.fc1.weight.numel():.2%}")

# 3. 对剪枝后的模型进行微调
# 微调时通常使用更小的学习率
fine_tune_optimizer = optim.Adam(model.parameters(), lr=0.0001)

# 微调循环（几个epoch即可）
num_fine_tune_epochs = 5
for epoch in range(num_fine_tune_epochs):
    model.train()
    running_loss = 0.0
    # 假设 train_loader 是你的数据加载器
    for inputs, labels in train_loader:
        fine_tune_optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()

        # **重要提示**：在反向传播时，被剪枝的权重（值为0）其梯度也将为0，不会被更新。
        fine_tune_optimizer.step()

        running_loss += loss.item()

    print(f'微调 Epoch [{epoch + 1}/{num_fine_tune_epochs}], 损失: {running_loss / len(train_loader):.4f}')

print("模型微调完成。")

# 4. 永久移除剪枝掩码
# 永久移除剪枝掩码，使剪枝效果固定下来
for module, name in parameters_to_prune:
    prune.remove(module, name)

# 此时，检查模型参数，被剪掉的权重已真正消失，模型大小减小。
print("剪枝掩码已永久移除。")
